"""
RRT_2D
@author: huiming zhou
"""

import os
import sys
import math
import numpy as np

sys.path.append(os.path.dirname(os.path.abspath(__file__)) +
                "/../../Sampling_based_Planning/")

#from Sampling_based_Planning.rrt_2D import env, plotting, utils
import env, plotting, utils


class Node:
    def __init__(self, n):
        self.x = n[0]
        self.y = n[1]
        self.parent = None


class Rrt:
    def __init__(self, s_start, s_goal, step_len, goal_sample_rate, iter_max):
        self.s_start = Node(s_start)
        self.s_goal = Node(s_goal)
        self.step_len = step_len   # 步长
        self.goal_sample_rate = goal_sample_rate # 采样点直接落在目标点上的概率
        self.iter_max = iter_max   # 迭代次数
        self.vertex = [self.s_start] # RRT tree

        self.env = env.Env()
        self.plotting = plotting.Plotting(s_start, s_goal)
        self.utils = utils.Utils()

        self.x_range = self.env.x_range  # x轴上的取值范围
        self.y_range = self.env.y_range  # y轴上的取值范围
        self.obs_circle = self.env.obs_circle
        self.obs_rectangle = self.env.obs_rectangle
        self.obs_boundary = self.env.obs_boundary

    def planning(self):
        # 采样（不超过迭代次数）
        for i in range(self.iter_max):
            node_rand = self.generate_random_node(self.goal_sample_rate) # 随机采样点
            node_near = self.nearest_neighbor(self.vertex, node_rand)    # RRT tree 上离采样点最近的一个点 
            node_new = self.new_state(node_near, node_rand)              # 在 step_len 允许范围内的新点

            # 如果新点和离采样点最近的一个点之间没有障碍物
            if node_new and not self.utils.is_collision(node_near, node_new):
                # 把新点加入到 RRT tree 中
                self.vertex.append(node_new)
                # 计算新点和目标的距离
                dist, _ = self.get_distance_and_angle(node_new, self.s_goal)

                # 如果新点和目标点的距离小于 step_len, 且之间没有障碍物，表示找到了一条可行路径
                # 为什么通过这种方式？因为采样点刚好落在目标点的概率非常低。
                if dist <= self.step_len and not self.utils.is_collision(node_new, self.s_goal):
                    self.new_state(node_new, self.s_goal)
                    return self.extract_path(node_new)

        return None

    def generate_random_node(self, goal_sample_rate):
        '''
        生成随机采样点
        '''
        # delta 变化量
        delta = self.utils.delta
        # np.random.random() 产生（0,1）的浮点数 
        # goal_sample_rate = 0.05
        # 表示0.95的概率是在随机方向上采样
        if np.random.random() > goal_sample_rate:
            # np.random.uniform    从一个均匀分布[low,high)中随机采样  左闭右开
            # self.x_range[0] = 0    self.x_range[1] = 50
            # self.y_range[0] = 0    self.y_range[1] = 30
            return Node((np.random.uniform(self.x_range[0] + delta, self.x_range[1] - delta),
                         np.random.uniform(self.y_range[0] + delta, self.y_range[1] - delta)))

        # 0.05的概率是落在目标点上    
        return self.s_goal

    @staticmethod
    def nearest_neighbor(node_list, n):
        '''
        计算RRT tree上离采样点最近的点
        '''
        # np.argmin 给出水平方向最小值的下标；
        return node_list[int(np.argmin([math.hypot(nd.x - n.x, nd.y - n.y)
                                        for nd in node_list]))]

    def new_state(self, node_start, node_end):
        '''
        生成新点
        '''
        # 计算随机采样点和离采样点最近的点的距离和方向
        dist, theta = self.get_distance_and_angle(node_start, node_end)

        dist = min(self.step_len, dist)
        # 计算出新点
        node_new = Node((node_start.x + dist * math.cos(theta),
                         node_start.y + dist * math.sin(theta)))
        node_new.parent = node_start

        return node_new

    def extract_path(self, node_end):
        '''
        根据父节点，追溯起点，形成可行路径
        '''
        path = [(self.s_goal.x, self.s_goal.y)]
        node_now = node_end

        while node_now.parent is not None:
            node_now = node_now.parent
            path.append((node_now.x, node_now.y))

        return path

    @staticmethod
    def get_distance_and_angle(node_start, node_end):
        dx = node_end.x - node_start.x
        dy = node_end.y - node_start.y
        return math.hypot(dx, dy), math.atan2(dy, dx)


def main():
    x_start = (2, 2)   # Starting node
    x_goal = (49, 24)  # Goal node

    rrt = Rrt(x_start, x_goal, 0.5, 0.05, 10000) #实例化
    path = rrt.planning()

    if path:
        rrt.plotting.animation(rrt.vertex, path, "RRT", True)
    else:
        print("No Path Found!")


if __name__ == '__main__':
    main()
